from dataclasses import dataclass
from typing import Optional
@dataclass
class MSEData:
    #-----输入数据----------#
    date: Optional[str] = None # 采集日期
    time: Optional[str] = None # 采集时间
    depth: Optional[float] = None  # 深度
    TQ_surf: Optional[float] = None # 顶驱扭矩
    RPM_surf: Optional[float] = None #顶驱转速
    vertical_pipe_press: Optional[float] = None #立管压力
    Q: Optional[float] = None  #入口流量
    ROP: Optional[float] = None # 钻时
    WOB: Optional[float] = None # 钻压

# 数据转换，把国际单位转为英制
def transform_obj(obj:MSEData)->MSEData:

    # 确保所有值都不是None再进行计算
    # 顶驱扭矩
    if obj.TQ_surf is not None:
        obj.TQ_surf = round(obj.TQ_surf * 1000 * 0.7375621, 2)  # KN-m to ft-lbs

    # 立管压力的转换
    if obj.vertical_pipe_press is not None:
        obj.vertical_pipe_press = round(obj.vertical_pipe_press / 0.06895, 2)  # KPa to psi

    # 入口流量的单位换算
    if obj.Q is not None:
        obj.Q = round(obj.Q / 3.7854, 2)  # L/min to 加仑/分钟

    # 钻时的转换
    if obj.ROP is not None:
        obj.ROP = round(obj.ROP * 3.28084, 2)  # m/hr to ft/hr

    # 钻压的转换
    if obj.WOB is not None:
        obj.WOB = round(obj.WOB * 9.8 * 1000 * 0.2248089, 2)  # KN to lbs
    return obj

# 四个计算公式
def calculate_mse(obj:MSEData, D=8.5, deltaP =0.0, Kn=0.0, Kt=0.0):

    mse_total = 4*safe_float(obj.WOB) / (3.14 * D**2) + 480*((obj.TQ_surf - Kt*deltaP)*obj.RPM_surf + Kt*deltaP*(obj.RPM_surf+Kn*obj.Q)) / (D**2 * obj.ROP)

    mse_downhole = 4*obj.WOB / (3.14*D**2) + 480*Kt*deltaP*(obj.RPM_surf+ Kn*obj.Q) / (D**2 * obj.ROP)

    DOC = (12/60)*obj.ROP / (obj.RPM_surf + Kn*obj.Q)

    DS = 2 * obj.WOB / (DOC * D)

    return mse_total, mse_downhole,DOC,DS

# 辅助函数：安全转换为浮点数
def safe_float(value):
    """安全地将值转换为float"""
    if value is None:
        return None
    try:
        return float(value)
    except (ValueError, TypeError):
        print(f"警告: 无法将值 '{value}' 转换为浮点数")
        return None